import cv2
import numpy as np
import time
from tqdm import tqdm
from matplotlib.colors import hsv_to_rgb
import imageio
import os
import shutil

path_gif_dir = "./output/{}/"

#              上、     右、   下、   左
xy_offset_4 = [[-1, 0], [0, 1], [1, 0], [0, -1]]
#              上、     右、   下、   左、    左上、   右上、   右下、 左下
xy_offset_8 = [[-1, 0], [0, 1], [1, 0], [0, -1], [-1, -1], [-1, 1], [1, 1], [1, -1]]
# 8邻域偏移量，也是判断前进方向的数组
# wallin的顺序有点说法，要和xy_offset_4对上
# 以右手定则来说，墙处于左上角或正上方时要先往左走，确保墙处于自己的右手边
# wallin的索引下标整除2即可得出xy_offset_4的索引下标，进而知道前进方向
#                  右上、  右、   右下、  下、   左下、   左、    左上、   上
wallin = np.array([[-1, 1], [0, 1], [1, 1], [1, 0], [1, -1], [0, -1], [-1, -1], [-1, 0]])


# 彩虹调色板
def rainbow_cmap(palette_size: int = 14):
    """
    """
    # 色相值
    hues = np.linspace(0, 1, palette_size)
    # 饱和度
    saturation = np.ones(palette_size)
    # 明度值
    value = np.ones(palette_size)
    # 将HSV颜色转换成RGB颜色
    colors_hsv = np.array([hues, saturation, value]).T  # .T为矩阵变形，换轴
    colors_rgb = hsv_to_rgb(colors_hsv)

    # 0-255
    color_map = np.clip(np.multiply(colors_rgb, 255), 0, 255)
    rainbow_colors = color_map.astype(np.uint8)
    return rainbow_colors


def get_edge_path(image, padding: int = 2):
    """
    获取能填充整个image的由外向内的环形路径，每条

    路径应能像漩涡一样由外向内走直至填满整个图像

    Parameters
    ----------
    :param image : np.ndarray
        二值图, shape=(h, w)
    :param padding : int
        Number of values padded to the edges of each axis

    Return
    ------
    :return : list[np.ndarray]
            返回n条完整的、能从外向内填充图像的路径
    """

    assert len(image.shape) == 2

    print(f"before padding of Image.shape{image.shape}")
    # 最外层填充2格像素黑边，方便后续计算
    padding_img = np.pad(image, padding, 'constant', constant_values=0)
    h, w = padding_img.shape
    print(f"after padding of Image.shape({h},{w})")

    # ------------------------------    # 第一步：右手法则获取边缘路径    # ------------------------------
    paths = get_path(padding_img)
    print(f"总共找到 {len(paths)} 条独立的边缘路径")
    path_save(padding_img, paths)
    path_point_save(padding_img, paths, "single_path_point_gif")

    # -----------------------------    # 第一步：联通不同的环路                # ------------------------------
    union_paths = link_path(padding_img, paths)
    path_point_save(padding_img, union_paths, "union_path_point_gif")


def get_path(image):
    """
    """
    h, w = image.shape
    img = np.copy(image)
    # 构建二维图像的一维索引
    indices = np.linspace(0, h * w - 1, h * w, dtype=np.int32).reshape((h, w))
    k = 1
    # 从第一层圈开始，这样不需要边界判断（最外一层视为0层
    min_bor = min(h, w)
    end_k = (min_bor >> 1) + (min_bor & 1)
    print(f"Image总圈数: {end_k}")

    paths = []
    # 一圈一圈由外向内遍历img
    pbar = tqdm(desc="尝试寻找起点")
    while k <= end_k:
        # 获取索引
        cur_indices = get_indices(indices, k)
        # 转成二维索引
        y = cur_indices // w
        x = cur_indices - y * w
        # 过滤非白色255像素
        white_mask = img[y, x] == 255
        y = y[white_mask]
        x = x[white_mask]
        for i in range(len(y)):
            row, col = y[i], x[i]
            # 判断是否为白色255像素，前面的边缘路径曾经有可能路过当前点
            if img[row, col] != 255:
                pbar.update(1)
                continue

            # 尝试获取前进方向（如果8邻域有墙存在
            way_face = get_way_face(img, row, col)
            if way_face > -1:
                # 以该点为起点，找到一条右手扶着墙走得路线重回该点
                point = np.array([row, col])
                path = right_hand_rule(img, point, point, way_face)
                # 将找到的路径的所有点置为1（表示已访问，只要不是255都表示已访问
                rows, cols = path[:, 0], path[:, 1]
                img[rows, cols] = 1  # 置为1是因为有一些特殊的图
                paths.append(path)

            pbar.update(1)
            # 向内缩一圈
        k += 1
    pbar.close()
    return paths


def get_indices(indices: np.ndarray, k: int):
    """
    获取由外向内环形索引数组

    indices : [
    [1,  2,  3,  4,  5],
    [6,  7,  8,  9, 10],
    [11,12, 13, 14, 15],
    [16,17, 18, 19, 20]]

    k = 0, result = [1, 2, 3, 4, 5, 6, 11, 16, 10, 15, 20, 17, 18, 19]

    k = 1, result = [7, 8, 9, 12, 14, 13]

    Parameters
    ----------
    :param indices : np.ndarray
            二维索引数组
    :param k : int
            第几层环型索引

    Return
    ------
    :return : np.ndarray
            返回数组的值列表
    """
    h, w = indices.shape
    # 顶行
    top_line = indices[k, k:w - k]
    # 左、右列（不包括首行
    left_column = indices[k + 1:h - k, k]
    right_column = indices[k + 1:h - k, w - k - 1]
    # 底行（不包括首、尾列
    bottom_line = indices[h - k - 1, k + 1:w - k - 1]
    return np.hstack((top_line, left_column, right_column, bottom_line))


def right_hand_rule(
        image,
        start_point: np.ndarray,
        end_point: np.ndarray,
        way_face: int) -> np.ndarray:
    """
    右手法则找到一条从起点start_point走到终点end_point的路

    start_point和end_point的像素值应该相等，其他与这两点像素

    值不等的视为墙

    Parameters
    ----------
    :param start_point : np.ndarray
            起点的坐标
    :param end_point : np.ndarray
            终点的坐标
    :param way_face : int
            初始前进方向，xy_offset_4的索引

    Return
    ------
    :return : np.ndarray
            返回一条从起点到终点所经过point的xy坐标列表
    """
    global xy_offset_4
    path = [start_point]
    rh_pbar = tqdm(desc=f"右手法则{start_point.tolist()} >>> {end_point.tolist()}", leave=False)
    # 初始化
    cur_point = start_point
    # 起点与终点是同一个点时，手动挪动一个点
    if np.all(cur_point == end_point):
        way_face, cur_point = next_point_by_right_hand_rule(image, cur_point, way_face)
        rh_pbar.update(1)
    # 找到终点则结束
    while not np.all(cur_point == end_point):
        path.append(cur_point)
        way_face, cur_point = next_point_by_right_hand_rule(image, cur_point, way_face)
        rh_pbar.update(1)
    rh_pbar.close()
    path.append(end_point)
    return np.array(path, dtype=np.int32)


def get_way_face(image, y, x) -> int:
    """    
    获取坐标[y, x]的前进方向，这里将非255值的像素视为墙
        
    Parameters    
    ----------    
    :param y : {int, np.int}    
    :param x : {int, np.int}        
    
    Return    
    ------    
    :return : int
            返回-1表示邻域内没有墙（黑色像素），否则返回xy_offset_4数组的索引下标
    """
    global xy_offset_4, wallin
    # 邻域子图
    sub_img = image[y - 1:y + 2, x - 1:x + 2]
    # 获取邻域非白色像素坐标（子图的
    # 取任意一堵墙（黑色像素，值=0）作为前进方向的参考
    # 优先取已经明确是属于其他边缘的像素作为墙（值=1的
    rows0, cols0 = np.where(sub_img == 0)
    rows1, cols1 = np.where(sub_img == 1)
    rows, cols = np.hstack((rows1, rows0)), np.hstack((cols1, cols0))
    if len(rows) == 0:  # 周围无墙
        return -1

    y2, x2 = rows[0] - 1, cols[0] - 1  # 将坐标转换到以[y, x]为原点
    # 计算前进方向
    wall_mask = np.all(np.array([y2, x2]) == wallin, axis=1)
    indices, = np.nonzero(wall_mask)
    index = indices[0] // 2
    # 这里整除2可以得出xy_offset_4的索引下标，即下一步的要走的点（即前进方向
    return index.item()


def next_point_by_right_hand_rule(image, point: np.ndarray, way_face: int) -> tuple[int, np.ndarray]:
    """
    根据当前步的前进方向，计算出下一步的xy坐标，并给出下一步的前进方向

    way_face是xy_offset_4的索引

    Parameters
    ----------
    :param point : np.ndarray
            当前步的xy坐标
    :param way_face : int
            当前步的前进方向

    Return
    ------
    :return : int, np.int
            下一步的前进方向, 下一步的xy坐标
    """
    global xy_offset_4  # 上、右、下、左
    # 假设当前预设的前进方向是左，那会先尝试往左走，
    # 如果左边是墙，根据`右手定则`会继续尝试往下走，
    # 如果下边也是墙，继续尝试往右走
    y1, x1 = point
    for i in range(4):
        y2, x2 = point + xy_offset_4[way_face]
        if image[y2, x2] != image[y1, x1]:
            way_face = (way_face + 3) & 3
            continue
        return (way_face + 1) & 3, np.array([y2, x2])
    return -1, point


def link_path(image, path_list: list[np.ndarray]):
    """
    将具有嵌套关系的环形路径联通

    Parameters
    ----------
    :param image : np.ndarray, shape(h, w)
    :param path_list : list[np.ndarray]
            环形路径列表

    Return
    ------
    :return : list[np.ndarray]
                返回联通后的环形漩涡路径
    """
    # 构建像素到路径的索引
    pixel_2_path_ref = np.zeros(image.shape, dtype=np.int32)
    for i, path in enumerate(path_list):
        y, x = path[:, 0], path[:, 1]
        pixel_2_path_ref[y, x] = i + 1  # 从1开始，0为黑色像素

    # ---------------------------
    # 独占式构建预联通队列
    # ---------------------------
    path_ref = dict()  # path所在queue_list索引的映射关系
    queue_list = []  # 联通的path队列
    ref_queue_list = []  # 联通的path索引
    for i, path in enumerate(path_list):
        y, x = path[0]
        parent_ref = np.min(pixel_2_path_ref[y - 1:y + 2, x - 1:x + 2])
        self_ref = pixel_2_path_ref[y, x]

        if parent_ref not in path_ref:
            path_ref[self_ref] = len(queue_list)
            queue_list.append([path])
            ref_queue_list.append([self_ref])
        else:
            queue_index = path_ref.pop(parent_ref)
            queue_list[queue_index].append(path)
            ref_queue_list[queue_index].append(self_ref)
            path_ref[self_ref] = queue_index

    union_path_list = []
    for i, queue in enumerate(queue_list):
        deeper_path = queue[-1]  # 最内层嵌套的环

        # 计算一个最远点
        y_mean, x_mean = np.mean(deeper_path[:, 0]), np.mean(deeper_path[:, 1])
        distances = np.linalg.norm(deeper_path - np.column_stack((y_mean, x_mean)), axis=1)
        dis_max_index = np.argmax(distances)
        anchor_point = deeper_path[dis_max_index]

        # 四个方向射出虚拟射线，保留命中最多嵌套环的
        keypoint_list = direct_line(anchor_point, pixel_2_path_ref, ref_queue_list[i][::-1])
        keypoint_list = keypoint_list[::-1]
        # 从最外圈开始，将原始路径从anchor处拆分并重组
        union_path = []
        for j in range(len(keypoint_list)):
            ref = ref_queue_list[i][j]
            _path = path_list[ref-1]
            if isinstance(keypoint_list[j], list):
                anchor_farthest = keypoint_list[j][-1]
                anchor_index, = np.where(np.all(_path == anchor_farthest, axis=1))
                anchor_index = anchor_index[0]
                seg_path = np.vstack((_path[anchor_index:], _path[:anchor_index], np.array(keypoint_list[j][::-1])))
            else:
                anchor = keypoint_list[j]
                anchor_index, = np.where(np.all(_path == anchor, axis=1))
                anchor_index = anchor_index[0]
                seg_path = np.vstack((_path[anchor_index:], _path[:anchor_index], np.array([keypoint_list[j]])))

            union_path.append(seg_path)

        union_path = np.concatenate(union_path, axis=0)
        union_path_list.append(union_path)

    return union_path_list


def direct_line(anchor_point, pixel_2_path_ref, ref_queue):
    max_list = []
    for xy_offset in xy_offset_4:
        _list = [anchor_point]
        while True:
            pre_point = _list[-1]
            if isinstance(_list[-1], list):
                pre_point = _list[-1][-1]

            cur_point = pre_point + xy_offset

            # 当前选中的线路覆盖了一个环的多个点
            if pixel_2_path_ref[cur_point[0], cur_point[1]] == pixel_2_path_ref[pre_point[0], pre_point[1]]:
                if isinstance(_list[-1], list):
                    _list[-1].append(cur_point)
                else:
                    _list.pop()
                    _list.append([pre_point, cur_point])
                continue

            if (len(_list) < len(ref_queue)
                    and pixel_2_path_ref[cur_point[0], cur_point[1]] == ref_queue[len(_list)]):
                _list.append(cur_point)
                continue

            break


        if len(_list) > len(max_list):
            max_list = _list

    return max_list


def path_save(image, path_list: list[np.ndarray]) -> None:
    """
    边缘路径gif保存
    """

    h, w = image.shape
    rgb_image = np.zeros(shape=(h, w, 3), dtype=np.uint8)
    rgb_image[~(image == 0)] = 255
    # 获取color map
    rainbow_palette = rainbow_cmap()

    global path_gif_dir
    out_path = os.path.normpath(os.path.join(path_gif_dir, './path_save.imageio.gif'))
    writer = imageio.get_writer(out_path, mode='I', fps=2, loop=0)

    writer.append_data(rgb_image)  # 先将未染色图像写入
    # 按路径顺序染色并写入gif文件
    for i, path in enumerate(tqdm(path_list, desc='path GIF SAVE')):
        color = rainbow_palette[i % len(rainbow_palette)]
        y, x = path[:, 0], path[:, 1]
        rgb_image[y, x] = color
        writer.append_data(rgb_image)
    writer.close()
    print('save ./path_save.imageio.gif done.')


def path_point_save(image, paths: list[np.ndarray], namespace: str) -> None:
    """
    路径点位gif保存
    :param image:
    :param paths:
    :param namespace:
    :return:
    """

    h, w = image.shape
    rgb_image = np.zeros(shape=(h, w, 3), dtype=np.uint8)
    rgb_image[~(image == 0)] = 255

    # 获取color map
    rainbow_palette = rainbow_cmap()
    global path_gif_dir
    output_folder = os.path.normpath(os.path.join(path_gif_dir, namespace))
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    for path_points in tqdm(paths, desc='路径点位 GIF SAVE'):
        order = len(os.listdir(output_folder)) + 1
        file_name = f'./path_points_save_{order:03d}.imageio.gif'
        out_path = os.path.normpath(os.path.join(output_folder, file_name))

        y, x = path_points[:, 0], path_points[:, 1]
        rgb_image[y, x] = 172  # 将路径初始化成灰色

        writer = imageio.get_writer(out_path, mode='I', fps=15, loop=0)
        writer.append_data(rgb_image)

        # 按路径点位顺序染色并写入gif文件
        for i, point in enumerate(tqdm(path_points, desc='path point GIF SAVE', leave=False)):
            color = rainbow_palette[i % len(rainbow_palette)]
            row, col = point[0], point[1]
            rgb_image[row, col] = color
            writer.append_data(rgb_image)
        writer.close()

        # 将当前路径所有点位染成深灰色，表示完成了
        rgb_image[y, x] = 50

    print(f'save all done.')


def get_filename_without_extension(path):
    # 提取完整的文件名（包括扩展名）
    filename_with_ext = os.path.basename(path)
    # 分离文件名和扩展名
    filename, ext = os.path.splitext(filename_with_ext)
    return filename


def main():
    image_path = "./test_img_7.png"
    image_name = get_filename_without_extension(image_path)
    global path_gif_dir
    path_gif_dir = path_gif_dir.format(image_name)
    if os.path.exists(path_gif_dir):
        shutil.rmtree(path_gif_dir)
    os.makedirs(path_gif_dir)

    # 以灰度模式读取图片
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    # 应用阈值进行二值化
    _, binary_image = cv2.threshold(image, 150, 255, cv2.THRESH_BINARY)
    # 获取边缘路径
    get_edge_path(binary_image)


# Backend tkagg is interactive backend. Turning interactive mode on.
if __name__ == '__main__':
    start_time = time.perf_counter()

    main()

    end_time = time.perf_counter()
    elapsed_time_ms = (end_time - start_time) * 1000
    print(f'耗时: {elapsed_time_ms:.3f}ms, done.')
